import socket
import pickle
import logging
import json
from typing import Dict, Any, Callable
from config import HOST, PORT
from cryptography.fernet import Fernet
from config import FERNET_KEY

logger = logging.getLogger(__name__)

class Client:
    def __init__(self, host: str = HOST, port: int = PORT):
        self.host = host
        self.port = port
        self.cipher_suite = Fernet(FERNET_KEY)
        
    def prepare_data(self, data: Dict[str, Any]) -> Dict[str, bytes]:
        """准备要发送的数据"""
        try:
            # 如果数据已经包含encrypted_data，直接使用
            if 'encrypted_data' in data:
                encrypted_data = data['encrypted_data']
            else:
                # 否则加密数据
                encrypted_data = self.cipher_suite.encrypt(
                    json.dumps(data).encode()
                )
            
            # 如果数据已经包含metadata_hash，直接使用
            if 'metadata_hash' in data:
                metadata_hash = data['metadata_hash']
            else:
                # 否则生成元数据哈希
                metadata_hash = self.cipher_suite.encrypt(
                    json.dumps(data.keys()).encode()
                )
            
            # 只返回服务器需要的字段
            return {
                'encrypted_data': encrypted_data,
                'metadata_hash': metadata_hash
            }
        except Exception as e:
            logger.error(f"准备数据时出错: {str(e)}")
            raise
            
    def send_data(self, data: Dict[str, Any]) -> None:
        """发送数据到服务器"""
        try:
            if not isinstance(data, dict):
                raise TypeError("数据必须是字典类型")
                
            # 准备数据
            prepared_data = self.prepare_data(data)
                
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(10)  # 设置超时时间
                s.connect((self.host, self.port))
                s.sendall(pickle.dumps(prepared_data))
                logger.info("数据成功发送到服务器")
        except socket.timeout:
            logger.error("连接超时")
            raise
        except ConnectionRefusedError:
            logger.error("连接被拒绝")
            raise
        except Exception as e:
            logger.error(f"发送数据失败: {str(e)}")
            raise

class Communicator:
    def __init__(self, host: str = HOST, port: int = PORT):
        self.host = host
        self.port = port
        
    def send_data(self, data: Dict[str, Any]) -> None:
        """发送数据到服务器"""
        try:
            if not isinstance(data, dict):
                raise TypeError("数据必须是字典类型")
                
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(10)  # 设置超时时间
                s.connect((self.host, self.port))
                s.sendall(pickle.dumps(data))
                logger.info("数据成功发送到服务器")
        except socket.timeout:
            logger.error("连接超时")
            raise
        except ConnectionRefusedError:
            logger.error("连接被拒绝")
            raise
        except Exception as e:
            logger.error(f"发送数据失败: {str(e)}")
            raise
            
    def start_server(self, callback: Callable):
        """启动服务器"""
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind((self.host, self.port))
                s.listen()
                logger.info(f"服务器启动在 {self.host}:{self.port}")
                
                while True:
                    conn, addr = s.accept()
                    with conn:
                        logger.info(f"连接来自 {addr}")
                        data = conn.recv(4096)
                        if not data:
                            break
                        try:
                            callback(pickle.loads(data))
                        except Exception as e:
                            logger.error(f"处理数据时出错: {str(e)}")
        except Exception as e:
            logger.error(f"服务器运行出错: {str(e)}")
            raise 